import cv2
import numpy as np
import os
import re
class MangaPanelExtractor:
    def __init__(self):
        self.white_threshold = 200  # 白色区域的阈值
        self.min_panel_size = 5000  # 最小分镜面积
        self.debug = True
        self.row_threshold = 150  # 同一行的最大y坐标差异

    def enhance_white_regions(self, image):
        """增强白色区域"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用自适应直方图均衡化
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        enhanced = clahe.apply(gray)
        
        # 使用高斯模糊减少噪点
        blurred = cv2.GaussianBlur(enhanced, (3, 3), 0)
        
        return blurred

    def find_white_regions(self, gray_image):
        """查找白色区域"""
        # 使用自适应阈值处理
        binary = cv2.adaptiveThreshold(
            gray_image,
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY,
            15,
            -2
        )
        
        # 创建一个较小的核用于形态学操作
        kernel = np.ones((2,2), np.uint8)
        
        # 进行开运算去除小的白点
        opened = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=1)
        
        # 进行闭运算填充小的黑点
        closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, kernel, iterations=1)
        
        return closed

    def find_panels(self, binary_image):
        """查找完整的分镜区域"""
        # 反转图像，使分镜区域为白色，边界为黑色
        inverted = cv2.bitwise_not(binary_image)
        
        # 使用连通组件分析找到所有可能的分镜
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(inverted, connectivity=8)
        
        # 过滤和处理分镜
        valid_panels = []
        for i in range(1, num_labels):  # 跳过背景（标签0）
            area = stats[i, cv2.CC_STAT_AREA]
            
            # 过滤太小的区域
            if area < self.min_panel_size:
                continue
                
            # 获取该区域的掩码
            mask = (labels == i).astype(np.uint8) * 255
            
            # 找到轮廓
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                # 使用最大的轮廓
                contour = max(contours, key=cv2.contourArea)
                valid_panels.append(contour)
        
        return valid_panels

    def extract_panel_with_contour(self, image, contour, padding=2):
        """使用轮廓提取分镜，并保持原始形状"""
        # 获取轮廓的边界框
        x, y, w, h = cv2.boundingRect(contour)
        
        # 添加padding
        x = max(0, x - padding)
        y = max(0, y - padding)
        w = min(image.shape[1] - x, w + 2 * padding)
        h = min(image.shape[0] - y, h + 2 * padding)
        
        # 直接提取ROI区域并返回，保持原始内容
        roi = image[y:y+h, x:x+w].copy()
        return roi, (x, y, w, h)

    def sort_panels(self, panels):
        """按照漫画阅读顺序（从左到右，从上到下）对分镜进行排序"""
        # 获取每个分镜的边界框
        panel_bounds = []
        for contour in panels:
            x, y, w, h = cv2.boundingRect(contour)
            # 存储轮廓和边界框信息
            panel_bounds.append({
                'contour': contour,
                'x': x,
                'y': y,
                'w': w,
                'h': h,
                'center_y': y + h/2  # 计算中心y坐标
            })
        
        # 按y坐标粗略分组为不同行
        rows = []
        current_row = []
        sorted_panels = sorted(panel_bounds, key=lambda p: p['center_y'])
        
        if sorted_panels:
            current_row = [sorted_panels[0]]
            last_y = sorted_panels[0]['center_y']
            
            for panel in sorted_panels[1:]:
                # 如果当前分镜与上一个分镜的y坐标差异小于阈值，认为是同一行
                if abs(panel['center_y'] - last_y) < self.row_threshold:
                    current_row.append(panel)
                else:
                    # 将当前行按x坐标排序
                    current_row.sort(key=lambda p: p['x'])
                    rows.append(current_row)
                    current_row = [panel]
                last_y = panel['center_y']
            
            # 添加最后一行
            if current_row:
                current_row.sort(key=lambda p: p['x'])
                rows.append(current_row)
        
        # 将所有行合并为一个列表，保持轮廓的顺序
        sorted_contours = []
        for row in rows:
            sorted_contours.extend([p['contour'] for p in row])
        
        return sorted_contours

    def extract_panels(self, image_path, output_dir):
        """提取漫画分镜"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"\n开始处理图片: {image_path}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 增强白色区域
        enhanced = self.enhance_white_regions(image)
        
        # 查找白色区域
        binary = self.find_white_regions(enhanced)
        
        # 保存二值化图像用于调试
        cv2.imwrite(os.path.join(output_dir, "binary.jpg"), binary)
        
        # 查找分镜
        panels = self.find_panels(binary)
        
        # 对分镜进行排序
        sorted_panels = self.sort_panels(panels)
        
        # 在调试图像上绘制分镜边界
        debug_image = image.copy()
        
        # 提取分镜
        panel_count = 0
        for contour in sorted_panels:
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            if area < self.min_panel_size:
                continue
            
            # 提取分镜
            panel, (x, y, w, h) = self.extract_panel_with_contour(image, contour)
            
            # 检查是否为空白区域
            gray_panel = cv2.cvtColor(panel, cv2.COLOR_BGR2GRAY)
            if np.mean(gray_panel) > self.white_threshold:
                continue
            
            # 检查非白色像素的比例
            non_white_ratio = np.sum(gray_panel < self.white_threshold) / (gray_panel.shape[0] * gray_panel.shape[1])
            if non_white_ratio < 0.3:  # 如果非白色像素太少，跳过
                continue
            
            # 在调试图像上绘制轮廓和编号
            cv2.drawContours(debug_image, [contour], -1, (0, 255, 0), 2)
            cv2.putText(debug_image, str(panel_count + 1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            panel_count += 1
            output_path = os.path.join(output_dir, f"panel_{panel_count:02d}.jpg")  # 使用两位数字编号
            cv2.imwrite(output_path, panel)
            print(f"保存分镜 {panel_count:02d}: {output_path}")
        
        # 保存调试图像
        cv2.imwrite(os.path.join(output_dir, "detected_panels.jpg"), debug_image)
        
        return panel_count

def natural_sort_key(s):
    """用于实现文件夹的自然排序"""
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', str(s))]

def main():
    # 创建实例
    extractor = MangaPanelExtractor()
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
    # 获取所有JPG文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg') or f.lower().endswith('.png')]
    
    print(f"\n找到 {len(image_files)} 个图片文件")
    
    if not image_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到图片文件！")
        return
    
    #进行自然排序
    image_files.sort(key=natural_sort_key)
    
    # 处理每个图片文件
    for image_file in image_files:
        print(f"\n处理文件: {image_file}")
        input_path = os.path.join(input_dir, image_file)
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(image_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {image_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {image_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 